import pygame
import sys
import time
import os
import cfg # Import cfg first
from modules import * # Then import modules which might use cfg

# --- Initialize Pygame and Mixer ---
pygame.init()
try:
    pygame.mixer.init()
    print("Pygame Mixer Initialized.")
    mixer_initialized = True
except pygame.error as e:
    print(f"错误: 初始化 Pygame Mixer 失败: {e}")
    mixer_initialized = False
    cfg.SOUND_ENABLED = False # Disable sound if mixer failed
    cfg.BGM_ENABLED = False   # Disable BGM if mixer failed

# --- Load Sounds (after mixer init attempt) ---
if mixer_initialized:
    try:
        all_sounds_loaded = True
        if os.path.exists(cfg.EAT_PATH):
             cfg.SFX_EAT = pygame.mixer.Sound(cfg.EAT_PATH)
        else:
             print(f"警告: 未找到音效文件: {cfg.EAT_PATH}")
             all_sounds_loaded = False

        if os.path.exists(cfg.DIE_PATH):
             cfg.SFX_DIE = pygame.mixer.Sound(cfg.DIE_PATH)
        else:
             print(f"警告: 未找到音效文件: {cfg.DIE_PATH}")
             all_sounds_loaded = False

        if os.path.exists(cfg.LEVELUP_PATH):
             cfg.SFX_LEVELUP = pygame.mixer.Sound(cfg.LEVELUP_PATH)
        else:
             print(f"警告: 未找到音效文件: {cfg.LEVELUP_PATH}")
             all_sounds_loaded = False

        if all_sounds_loaded and cfg.SFX_EAT and cfg.SFX_DIE and cfg.SFX_LEVELUP:
            cfg.SFX_EAT.set_volume(0.7)
            cfg.SFX_DIE.set_volume(0.8)
            cfg.SFX_LEVELUP.set_volume(0.8)
            cfg.SOUND_ENABLED = True
            print("音效加载成功！")
        else:
            cfg.SOUND_ENABLED = False
            # Ensure specific sounds are None if they failed to load
            if not os.path.exists(cfg.EAT_PATH): cfg.SFX_EAT = None
            if not os.path.exists(cfg.DIE_PATH): cfg.SFX_DIE = None
            if not os.path.exists(cfg.LEVELUP_PATH): cfg.SFX_LEVELUP = None
            print("警告: 部分或全部音效加载失败，音效将被禁用。")

    except pygame.error as e:
        print(f"警告: 加载音效时发生错误: {e}. 音效将被禁用。")
        cfg.SOUND_ENABLED = False
        cfg.SFX_EAT = None
        cfg.SFX_DIE = None
        cfg.SFX_LEVELUP = None
else:
     print("警告: Mixer 未初始化，无法加载音效。")
     cfg.SOUND_ENABLED = False
     cfg.SFX_EAT = None
     cfg.SFX_DIE = None
     cfg.SFX_LEVELUP = None

# --- Load Music (after mixer init attempt) ---
bgm_loaded_successfully = False
if cfg.BGM_ENABLED and mixer_initialized:
    try:
        pygame.mixer.music.load(cfg.BGM_PATH)
        pygame.mixer.music.set_volume(0.5) # Set volume
        bgm_loaded_successfully = True
        print("背景音乐加载成功。")
    except pygame.error as e:
        print(f"错误: 加载背景音乐失败: {e}")
        cfg.BGM_ENABLED = False # Disable BGM if loading failed
else:
    if cfg.BGM_ENABLED and not mixer_initialized:
        print("警告: Mixer 未初始化，无法加载背景音乐。")
    # Keep BGM_ENABLED as False if file wasn't found initially in cfg.py
    cfg.BGM_ENABLED = False

session_high_score = 0

def load_achievements():
    """加载成就排行榜"""
    if not os.path.exists(cfg.ACHIEVEMENTS_FILE):
        return {}
    with open(cfg.ACHIEVEMENTS_FILE, 'r') as file:
        lines = file.readlines()
        achievements = {}
        for line in lines:
            mode, score = line.strip().split(',')
            achievements[mode] = int(score)
        return achievements

def save_achievements(achievements):
    """保存成就排行榜"""
    with open(cfg.ACHIEVEMENTS_FILE, 'w') as file:
        for mode, score in achievements.items():
            file.write(f"{mode},{score}\n")

def game_loop(screen, clock, cfg, snake_color, mode, difficulty_name):
    global session_high_score

    game_over = False
    paused = False
    current_level = 0
    score = 0
    target_score = 0
    base_speed = 15
    has_poison = False # Flag to track if poison is active in the current level/mode
    traps = []
    difficulty_multiplier = cfg.DIFFICULTY_SETTINGS.get(difficulty_name, 1.0)

    # --- Initialize game elements ---
    snake = Snake(cfg, color=snake_color)
    food_manager = FoodManager(cfg)
    # Initial food spawn (no traps or poison initially)
    food_manager.spawn_food(snake.coords, traps)

    # --- Mode specific setup ---
    if mode == "关卡模式":
        current_level = 1 # Start at level 1
    elif mode == "传统模式":
        current_level = 0 # Use 0 to indicate traditional mode
        base_speed = 15 # Define base speed for traditional mode
        difficulty_multiplier = cfg.DIFFICULTY_SETTINGS.get(difficulty_name, 1.0)
        print(f"进入传统模式 (难度: {difficulty_name}, 速度乘数: {difficulty_multiplier:.1f})")
        # Note: Traditional mode starts with no traps or poison unless added later

    # --- Level setup function ---
    def setup_level(level_num):
        nonlocal target_score, base_speed, has_poison, traps, food_manager, snake, current_speed
        if level_num > 0 and level_num <= cfg.MAX_LEVEL: # Check valid level number
            level_settings = cfg.LEVELS[level_num - 1]
            target_score = level_settings.get("target_score", 0) # Default 0 if missing
            base_speed = level_settings.get("speed", 15) # Default 15 if missing
            has_poison = level_settings.get("poison", False) # Default False if missing
            trap_count = level_settings.get("traps", 0) # Default 0 if missing

            # Reset snake for the new level
            snake = Snake(cfg, color=snake_color)
            # Place traps for the level, avoiding the new snake's position
            traps[:] = place_traps(cfg, trap_count, snake.coords)
            # Clear any existing poison and spawn initial food/poison for the level
            food_manager.clear_poison()
            food_manager.spawn_food(snake.coords, traps)
            if has_poison:
                # print(f"关卡 {level_num}: 毒药已启用!") # Notify when poison starts (optional)
                food_manager.spawn_poison(snake.coords, traps)

            # Calculate the actual speed for this level and difficulty
            current_speed = base_speed * difficulty_multiplier
            print(f"设置关卡 {level_num}: 目标 {target_score}, 基础速度 {base_speed}, 毒 {has_poison}, 陷阱 {len(traps)}, 实际速度 {current_speed:.1f}")
            return True
        else:
            print(f"错误: 无效的关卡号 {level_num}")
            return False # Invalid level number

    # --- Initial setup based on mode ---
    if mode == "关卡模式":
        if not setup_level(current_level):
             print("错误：无法设置初始关卡！")
             return "MENU", score # Return to menu if setup fails
    else: # Traditional mode initial speed calculation
        current_speed = base_speed * difficulty_multiplier
        traps[:] = [] # Ensure no traps in traditional mode start
        food_manager.clear_poison() # Ensure no poison in traditional mode start
        has_poison = False # Explicitly set for traditional mode start
        print(f"传统模式初始速度: {current_speed:.1f}")

    # --- Game Loop ---
    while not game_over:
        # --- Event Handling ---
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.stop()
                pygame.quit()
                sys.exit()
            elif event.type == pygame.KEYDOWN:
                if not paused: # Process movement keys only if not paused
                    if event.key == pygame.K_UP: snake.change_direction("UP")
                    elif event.key == pygame.K_DOWN: snake.change_direction("DOWN")
                    elif event.key == pygame.K_LEFT: snake.change_direction("LEFT")
                    elif event.key == pygame.K_RIGHT: snake.change_direction("RIGHT")
                # Pause/Resume (Space) and Exit (Esc) work whether paused or not
                if event.key == pygame.K_SPACE:
                    paused = not paused
                    if paused:
                         if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.pause()
                         display_pause_message(screen, cfg) # Show pause message immediately
                    else:
                         if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.unpause()
                         # No need to explicitly clear pause message, screen redraw will do it
                elif event.key == pygame.K_ESCAPE:
                     if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.fadeout(500) # Fade out BGM
                     return "MENU", score # Return to menu

        # If paused, skip game updates and redraw pause message
        if paused:
            # display_pause_message(screen, cfg) # Redraw needed if background changes
            pygame.display.flip() # Update display to show message overlay
            clock.tick(10) # Lower tick rate during pause
            continue # Skip rest of the game loop

        # --- Game Logic Update ---
        snake.update_direction()
        ate_food = False
        ate_poison = False
        grow_snake = False
        shrink_snake = False
        head = snake.head_coord

        if head: # Only process if snake head exists
            # Check collision with RED food
            if head == food_manager.get_food_pos():
                ate_food = True
                grow_snake = True
                score += 10
                if cfg.SOUND_ENABLED and cfg.SFX_EAT: cfg.SFX_EAT.play()

            # Check collision with BLUE poison (only if poison is active)
            poison_pos = food_manager.get_poison_pos()
            if has_poison and poison_pos and head == poison_pos:
                 ate_poison = True
                 shrink_snake = True
                 score -= 5 # Penalty for eating poison
                 if score < 0: score = 0 # Score doesn't go below zero
                 # Use the same 'eat' sound for poison? Or add a different one?
                 if cfg.SOUND_ENABLED and cfg.SFX_EAT: cfg.SFX_EAT.play() # Current uses eat sound

        # Move the snake based on flags
        snake.move(grow=grow_snake, shrink=shrink_snake)

        # Spawn new items if eaten
        if ate_food: food_manager.spawn_food(snake.coords, traps)
        # Only respawn poison if it was eaten and poison is active
        if ate_poison and has_poison: food_manager.spawn_poison(snake.coords, traps)

        # Check for game over conditions (wall, self, traps)
        if snake.check_collision(traps, mode="wall"):
            if cfg.SOUND_ENABLED and cfg.SFX_DIE: cfg.SFX_DIE.play()
            game_over = True
            continue # Skip rest of loop if game over

        # --- Level Mode Progression ---
        if mode == "关卡模式" and score >= target_score:
            # ---- MODIFICATION HERE ----
            if cfg.SOUND_ENABLED and cfg.SFX_LEVELUP:
                # Play level up sound for max 1500 (1.5 seconds)
                cfg.SFX_LEVELUP.play(maxtime=1500)
            # ---- END MODIFICATION ----

            current_level += 1
            if current_level > cfg.MAX_LEVEL:
                 print("恭喜通关所有关卡!")
                 game_over = True # Completed all levels
                 continue
            else:
                display_level_transition(screen, cfg, current_level)
                if not setup_level(current_level): # Setup the next level
                    print(f"错误：无法设置关卡 {current_level}!")
                    if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.fadeout(500)
                    return "MENU", score # Return to menu if setup fails
                score = 0 # Reset score for the new level
                # Speed is set within setup_level

        # --- Rendering ---
        screen.fill(cfg.BLACK) # Clear screen
        drawGameGrid(cfg, screen) # Draw grid lines
        if traps: draw_traps(screen, traps, cfg) # Draw traps (under snake/food)
        snake.draw(screen) # Draw the snake
        food_manager.draw(screen) # Draw food and poison
        # Display score/level info (pass 0 for target if traditional)
        display_score(screen, score, current_level, target_score if mode == "关卡模式" else 0, cfg)

        pygame.display.flip() # Update the full screen
        clock.tick(current_speed) # Control game speed

    # --- Game Over Handling ---
    if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.fadeout(1000) # Fade out music slowly
    if score > session_high_score: session_high_score = score # Update high score

    # Optional delay before showing end screen
    # pygame.time.wait(500)

    # Show end interface and get user action
    end_action = end_interface(screen, cfg, score, session_high_score, mode, current_level if mode=="关卡模式" else None)

    # Ensure music is fully stopped before potentially restarting or going to menu
    if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.stop()

    return end_action, score

def main():
    global session_high_score
    # Pygame, Mixer init, Sound/Music loading happens at the top now

    screen = pygame.display.set_mode(cfg.SCREENSIZE)
    pygame.display.set_caption("贪吃蛇 大作战 (音乐版)")
    clock = pygame.time.Clock()

    # Set initial game state variables
    current_snake_color = cfg.GREEN # Default color
    current_mode = "关卡模式" # Default mode
    current_difficulty = cfg.DEFAULT_DIFFICULTY # Default difficulty

    # Load achievements
    achievements = load_achievements()

    # --- Main Application Loop ---
    while True:
        # Stop any music before showing the menu
        if cfg.BGM_ENABLED and mixer_initialized: pygame.mixer.music.stop()

        # Show start menu and get user choices
        menu_result = start_menu(screen, cfg, current_snake_color, current_difficulty)

        if menu_result is None: # User quit from menu (e.g., closed window, clicked Quit)
             break # Exit the main loop

        # Unpack menu results
        selected_color, selected_mode, selected_difficulty = menu_result

        # Update current settings for the game loop and potential restarts
        current_snake_color = selected_color
        current_mode = selected_mode
        current_difficulty = selected_difficulty

        # Play BGM before starting the game loop if enabled and loaded
        if cfg.BGM_ENABLED and mixer_initialized and bgm_loaded_successfully:
            pygame.mixer.music.play(-1) # -1 loops indefinitely

        # --- Run the game ---
        action, final_score = game_loop(screen, clock, cfg, current_snake_color, current_mode, current_difficulty)

        # Update achievements
        if current_mode in achievements:
            if final_score > achievements[current_mode]:
                achievements[current_mode] = final_score
        else:
            achievements[current_mode] = final_score

        # Save achievements
        save_achievements(achievements)

        # --- Handle post-game action ---
        if action == "QUIT":
            break # Exit the main loop
        elif action == "MENU":
            continue # Go back to the start of the main loop to show the menu again
        elif action == "RESTART":
             # Keep restarting until the user chooses Menu or Quit from the end screen
             while action == "RESTART":
                 print(f"重新开始 {current_mode} (难度: {current_difficulty})...")
                 # Restart BGM if enabled
                 if cfg.BGM_ENABLED and mixer_initialized and bgm_loaded_successfully:
                      pygame.mixer.music.play(-1)
                 # Run the game loop again with the exact same settings
                 action, final_score = game_loop(screen, clock, cfg, current_snake_color, current_mode, current_difficulty)
                 # Update achievements
                 if current_mode in achievements:
                     if final_score > achievements[current_mode]:
                         achievements[current_mode] = final_score
                 else:
                     achievements[current_mode] = final_score
                 # Save achievements
                 save_achievements(achievements)
             # After the restart loop, check the action again
             if action == "QUIT":
                 break # Exit main loop if Quit was chosen after a restart
             # If action is "MENU", the outer loop's continue will handle it

    # --- Cleanup ---
    pygame.quit()
    sys.exit()

if __name__ == "__main__":
    main()
